// pvpkeet/pvpkeet.go
package pvpkeet

import (
	"crypto/sha256"

	"github.com/Nik-U/pbc"
)

type PVPKEET struct {
	Params *SystemParams
}

// H1 是协议中的哈希函数: {0,1}* × Z_q → Z_q
func (p *PVPKEET) H1(m []byte, t *pbc.Element) *pbc.Element {
	var seed []byte
	if t != nil {
		seed = append(m, t.Bytes()...)
	} else {
		seed = m
	}
	// 将哈希结果映射到Zr
	arr := sha256.Sum256(seed)
	return p.Params.Pairing.NewZr().SetFromHash(arr[:])
}

// Setup 初始化系统参数
func Setup(curveType string) *PVPKEET {
	// 使用 Type A 配对，对应 Python Charm-Crypto 的 SS (Supersingular) 曲线
	params := pbc.GenerateA(160, 512)
	pairing := params.NewPairing()

	sysParams := &SystemParams{
		Pairing: pairing,
		G:       pairing.NewG1().Rand(),
		H:       pairing.NewG1().Rand(),
	}
	return &PVPKEET{Params: sysParams}
}

// UserKeyGen 用户密钥生成
func (p *PVPKEET) UserKeyGen() *UserKeyPair {
	s := p.Params.Pairing.NewZr().Rand()
	u := p.Params.Pairing.NewG1().PowZn(p.Params.G, s)
	return &UserKeyPair{Pku: u, Sku: s}
}

// CloudKeyGen 云服务器密钥生成
func (p *PVPKEET) CloudKeyGen() *CloudKeyPair {
	alpha := p.Params.Pairing.NewZr().Rand()
	pkc := p.Params.Pairing.NewG1().PowZn(p.Params.H, alpha)
	return &CloudKeyPair{Pkc: pkc, Skc: alpha}
}

// Encryption 加密算法
func (p *PVPKEET) Encryption(pku, pkc *pbc.Element, m []byte) *Ciphertext {
	pairing := p.Params.Pairing
	t := pairing.NewZr().Rand()
	k := pairing.NewZr().Rand()
	r := p.H1(m, t)

	// group_m = group.hash(m, G1)
	arr := sha256.Sum256(m)
	group_m := pairing.NewG1().SetFromHash(arr[:])

	c1 := pairing.NewG1().PowZn(pku, r)
	c2_base := pairing.NewG1().PowZn(p.Params.G, r)
	c2 := pairing.NewG1().Mul(c2_base, group_m)
	c3 := k

	pkc_r := pairing.NewG1().PowZn(pkc, r)
	h_kr := pairing.NewG1().PowZn(p.Params.H, pairing.NewZr().Mul(k, r))
	c4 := pairing.NewG1().Mul(pkc_r, h_kr)

	return &Ciphertext{C1: c1, C2: c2, C3: c3, C4: c4}
}

// Test 测试两个密文是否对应相同明文
func (p *PVPKEET) Test(CTA, CTB *Ciphertext, skc_a, skc_b *pbc.Element) bool {
	pairing := p.Params.Pairing

	// EA = CTA.c4 ^ (1 / (skc_a + CTA.c3))
	exp_a := pairing.NewZr().Add(skc_a, CTA.C3)
	exp_a.Invert(exp_a)
	EA := pairing.NewG1().PowZn(CTA.C4, exp_a)

	// EB = CTB.c4 ^ (1 / (skc_b + CTB.c3))
	exp_b := pairing.NewZr().Add(skc_b, CTB.C3)
	exp_b.Invert(exp_b)
	EB := pairing.NewG1().PowZn(CTB.C4, exp_b)

	// lhs = e(EA / EB, g)
	base_lhs := pairing.NewG1().Div(EA, EB)
	lhs := pairing.NewGT().Pair(base_lhs, p.Params.G)

	// rhs = e(h, CTA.c2 / CTB.c2)
	base_rhs := pairing.NewG1().Div(CTA.C2, CTB.C2)
	rhs := pairing.NewGT().Pair(p.Params.H, base_rhs)

	return lhs.Equals(rhs)
}

// Proof 生成验证证明
func (p *PVPKEET) Proof(CTA, CTB *Ciphertext, skc_a, skc_b, pku *pbc.Element) *Proof {
	pairing := p.Params.Pairing
	t := pairing.NewZr().Rand()
	t_inv := pairing.NewZr().Invert(t)

	// pi_A = CTA.c4 ^ (t / (skc_a + CTA.c3))
	exp_a_base := pairing.NewZr().Add(skc_a, CTA.C3)
	exp_a := pairing.NewZr().Mul(t, exp_a_base.Invert(exp_a_base))
	pi_A := pairing.NewG1().PowZn(CTA.C4, exp_a)

	// pi_B = CTB.c4 ^ (t / (skc_b + CTB.c3))
	exp_b_base := pairing.NewZr().Add(skc_b, CTB.C3)
	exp_b := pairing.NewZr().Mul(t, exp_b_base.Invert(exp_b_base))
	pi_B := pairing.NewG1().PowZn(CTB.C4, exp_b)

	theta := pairing.NewG1().PowZn(p.Params.H, t)
	theta_v := pairing.NewG1().PowZn(pku, t_inv)

	return &Proof{Pi_A: pi_A, Pi_B: pi_B, Theta: theta, Theta_v: theta_v}
}

// Verify 验证算法
func (p *PVPKEET) Verify(proof *Proof, CTA, CTB *Ciphertext, pkc, pku *pbc.Element) bool {
	pairing := p.Params.Pairing

	// Verify(1)
	v1_base_lhs := pairing.NewG1().Div(proof.Pi_A, proof.Pi_B)
	v1_lhs := pairing.NewGT().Pair(v1_base_lhs, p.Params.G)
	v1_base_rhs := pairing.NewG1().Div(CTA.C2, CTB.C2)
	v1_rhs := pairing.NewGT().Pair(proof.Theta, v1_base_rhs)
	v1 := v1_lhs.Equals(v1_rhs)
	if !v1 {
		return false
	}

	// Verify(2)
	v2_lhs := pairing.NewGT().Pair(CTA.C1, p.Params.H)
	v2_rhs := pairing.NewGT().Pair(proof.Pi_A, proof.Theta_v)
	v2 := v2_lhs.Equals(v2_rhs)
	if !v2 {
		return false
	}

	// Verify(3)
	v3_base_rhs := pairing.NewG1().PowZn(p.Params.H, CTA.C3)
	v3_base_rhs.Mul(v3_base_rhs, pkc)
	v3_lhs := pairing.NewGT().Pair(proof.Pi_A, v3_base_rhs)
	v3_rhs := pairing.NewGT().Pair(CTA.C4, proof.Theta)
	v3 := v3_lhs.Equals(v3_rhs)
	if !v3 {
		return false
	}

	// Verify(4)
	// ( reusing v3_base_rhs as term )
	v4_lhs := pairing.NewGT().Pair(CTA.C4, pku)
	v4_rhs := pairing.NewGT().Pair(CTA.C1, v3_base_rhs)
	v4 := v4_lhs.Equals(v4_rhs)
	if !v4 {
		return false
	}

	// Verify(5)
	v5_lhs := pairing.NewGT().Pair(proof.Theta, proof.Theta_v)
	v5_rhs := pairing.NewGT().Pair(p.Params.H, pku)
	v5 := v5_lhs.Equals(v5_rhs)
	if !v5 {
		return false
	}

	return true
}
